import sys
import psutil
import time
import platform
import subprocess
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QListWidget, QLineEdit, QPushButton, QTextEdit,
                             QComboBox, QMessageBox, QInputDialog, QDialog, QDialogButtonBox,
                             QRadioButton, QButtonGroup, QAbstractItemView)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt6.QtGui import QTextCursor, QIntValidator

class MonitorThread(QThread):
    process_found = pyqtSignal(int, str)  # 添加进程名称参数
    update_status = pyqtSignal(str)

    def __init__(self, target, monitor_type):
        super().__init__()
        self.target = target
        self.monitor_type = monitor_type  # 'pid' 或 'name'
        self.running = True

    def run(self):
        while self.running:
            for proc in psutil.process_iter(['pid', 'name']):
                if not self.running:
                    return
                try:
                    if self.monitor_type == 'name' and proc.info['name'] == self.target:
                        self.process_found.emit(proc.info['pid'], proc.info['name'])
                    elif self.monitor_type == 'pid' and str(proc.info['pid']) == self.target:
                        self.process_found.emit(proc.info['pid'], proc.info['name'])
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            time.sleep(0.1)

    def stop(self):
        self.running = False
class PasswordDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("密码验证")
        self.setFixedSize(300, 120)
        
        layout = QVBoxLayout()
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.EchoMode.Password)
        self.password_input.setPlaceholderText("输入管理员密码")
        
        button_box = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        
        layout.addWidget(QLabel("检测到受限进程，请输入密码:"))
        layout.addWidget(self.password_input)
        layout.addWidget(button_box)
        self.setLayout(layout)
class SystemMonitor(QMainWindow):
    def __init__(self):
        super().__init__()
        self.password = ""
        self.target = ""
        self.monitor_type = "name"  # 默认监控类型
        self.monitor_thread = None
        self.cmd_allowed = True
        self.password_verified = False
        self.init_ui()
        self.refresh_process_list()
        self.scan_disks()

    def init_ui(self):
        self.setWindowTitle("程序监控器")
        self.setGeometry(100, 100, 1200, 600)

        main_widget = QWidget()
        main_layout = QHBoxLayout()
        
        # 左侧面板
        left_panel = QWidget()
        left_layout = QVBoxLayout()
        
        # 监控类型选择
        self.monitor_type_group = QButtonGroup(self)
        self.name_radio = QRadioButton("按进程名称监控")
        self.pid_radio = QRadioButton("按PID监控")
        self.name_radio.setChecked(True)
        self.monitor_type_group.addButton(self.name_radio)
        self.monitor_type_group.addButton(self.pid_radio)
        
        type_layout = QHBoxLayout()
        type_layout.addWidget(self.name_radio)
        type_layout.addWidget(self.pid_radio)
        left_layout.addLayout(type_layout)

        # 进程列表
        self.process_list = QListWidget()
        self.process_list.setSelectionMode(QAbstractItemView.SelectionMode.SingleSelection)
        self.process_list.itemDoubleClicked.connect(self.select_from_list)
        left_layout.addWidget(QLabel("进程列表（双击选择）:"))
        left_layout.addWidget(self.process_list)

        # 目标输入
        self.target_entry = QLineEdit()
        self.target_entry.setPlaceholderText("输入进程名称或选择下方进程")
        self.pid_validator = QIntValidator(1, 65535, self)
        left_layout.addWidget(QLabel("监控目标:"))
        left_layout.addWidget(self.target_entry)

        # 密码设置
        self.password_entry = QLineEdit()
        self.password_entry.setPlaceholderText("输入管理密码")
        self.password_entry.setEchoMode(QLineEdit.EchoMode.Password)
        left_layout.addWidget(QLabel("设置密码:"))
        left_layout.addWidget(self.password_entry)

        # 功能按钮
        self.set_password_btn = QPushButton("设置/更新密码")
        self.set_password_btn.clicked.connect(self.save_password)
        self.start_btn = QPushButton("开始监控")
        self.start_btn.clicked.connect(self.start_monitoring)
        self.stop_btn = QPushButton("停止监控")
        self.stop_btn.clicked.connect(self.stop_monitoring)
        self.stop_btn.setEnabled(False)

        left_layout.addWidget(self.set_password_btn)
        left_layout.addWidget(self.start_btn)
        left_layout.addWidget(self.stop_btn)

        left_panel.setLayout(left_layout)

        # 右侧面板
        right_panel = QWidget()
        right_layout = QVBoxLayout()
        
        # 输出信息
        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        self.output_text.setStyleSheet("font-family: Consolas; background-color: #1e1e1e; color: #dcdcdc;")
        right_layout.addWidget(QLabel("监控日志:"))
        right_layout.addWidget(self.output_text)

        # 命令执行
        self.disk_combo = QComboBox()
        self.cmd_entry = QLineEdit()
        self.cmd_entry.setPlaceholderText("输入系统命令")
        self.cmd_btn = QPushButton("执行命令")
        self.cmd_btn.clicked.connect(self.execute_cmd)
        
        cmd_layout = QHBoxLayout()
        cmd_layout.addWidget(self.disk_combo, 1)
        cmd_layout.addWidget(self.cmd_entry, 3)
        cmd_layout.addWidget(self.cmd_btn, 1)
        
        right_layout.addWidget(QLabel("命令执行（选择磁盘）:"))
        right_layout.addLayout(cmd_layout)

        right_panel.setLayout(right_layout)

        main_layout.addWidget(left_panel, 1)
        main_layout.addWidget(right_panel, 2)
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        # 信号连接
        self.name_radio.toggled.connect(self.update_monitor_type)
        self.pid_radio.toggled.connect(self.update_monitor_type)
        QTimer.singleShot(0, self.update_input_validation)

    def update_monitor_type(self):
        """更新监控类型并刷新输入验证"""
        self.monitor_type = "pid" if self.pid_radio.isChecked() else "name"
        self.target_entry.clear()
        self.update_input_validation()
        self.refresh_process_list()

    def update_input_validation(self):
        """根据监控类型更新输入验证"""
        if self.monitor_type == "pid":
            self.target_entry.setValidator(self.pid_validator)
            self.target_entry.setPlaceholderText("输入进程PID（数字）")
            self.process_list.setEnabled(False)
        else:
            self.target_entry.setValidator(None)
            self.target_entry.setPlaceholderText("输入进程名称或双击下方进程")
            self.process_list.setEnabled(True)

    def refresh_process_list(self):
        """刷新进程列表，根据监控类型显示不同信息"""
        self.process_list.clear()
        try:
            for proc in psutil.process_iter(['pid', 'name']):
                if self.monitor_type == "name":
                    item = f"{proc.info['name']} (PID: {proc.info['pid']})"
                else:
                    item = f"PID: {proc.info['pid']} - {proc.info['name']}"
                self.process_list.addItem(item)
        except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
            self.append_output(f"刷新进程列表错误: {str(e)}")

    def select_from_list(self, item):
        """从列表选择目标"""
        text = item.text()
        if self.monitor_type == "name":
            self.target_entry.setText(text.split(' (')[0])
        else:
            self.target_entry.setText(text.split(' ')[1])

    def save_password(self):
        """保存密码"""
        password = self.password_entry.text()
        if len(password) >= 4:
            self.password = password
            self.append_output("密码已更新")
        else:
            QMessageBox.critical(self, "错误", "密码长度至少4位")

    def validate_target(self):
        """验证目标输入有效性"""
        target = self.target_entry.text()
        if not target:
            QMessageBox.critical(self, "错误", "请输入监控目标")
            return False
        
        if self.monitor_type == "pid":
            if not target.isdigit():
                QMessageBox.critical(self, "错误", "PID必须是数字")
                return False
            try:
                psutil.Process(int(target))
            except psutil.NoSuchProcess:
                QMessageBox.critical(self, "错误", "指定PID的进程不存在")
                return False
        return True

    def start_monitoring(self):
        """启动监控"""
        if not self.validate_target():
            return
        
        self.target = self.target_entry.text()
        
        if self.target.lower() == "cmd.exe" and self.monitor_type == "name":
            reply = QMessageBox.question(
                self, "警告",
                "监控cmd.exe可能导致功能异常，是否继续？",
                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
            )
            if reply == QMessageBox.StandardButton.No:
                return
            self.cmd_allowed = False
            self.cmd_entry.setEnabled(False)
        
        self.monitor_thread = MonitorThread(self.target, self.monitor_type)
        self.monitor_thread.process_found.connect(self.handle_process)
        self.monitor_thread.start()
        
        self.toggle_controls(True)
        self.append_output(f"开始监控 [{'PID' if self.monitor_type == 'pid' else '名称'}]：{self.target}")

    def stop_monitoring(self):
        """停止监控"""
        if self.monitor_thread:
            self.monitor_thread.stop()
            self.monitor_thread = None
        self.toggle_controls(False)
        self.append_output("监控已停止")

    def toggle_controls(self, monitoring):
        """切换控件状态"""
        self.start_btn.setEnabled(not monitoring)
        self.stop_btn.setEnabled(monitoring)
        self.set_password_btn.setEnabled(not monitoring)
        self.password_entry.setEnabled(not monitoring)
        self.target_entry.setEnabled(not monitoring)
        self.cmd_entry.setEnabled(self.cmd_allowed)

    def handle_process(self, pid, name):
        """处理发现的目标进程"""
        if not self.password_verified:
            self.terminate_process(pid)
            self.ask_password(pid, name)
        else:
            self.resume_process(pid)

    def ask_password(self, pid, name):
        """请求密码验证"""
        dialog = PasswordDialog(self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            if dialog.password_input.text() == self.password:
                self.password_verified = True
                self.append_output(f"已授权进程：{name} (PID: {pid})")
                self.resume_process(pid)
            else:
                QMessageBox.critical(self, "错误", "密码验证失败")
                self.password_verified = False

    def terminate_process(self, pid):
        """终止进程"""
        try:
            psutil.Process(pid).terminate()
            self.append_output(f"已拦截未授权进程 (PID: {pid})")
        except psutil.NoSuchProcess:
            self.append_output(f"进程已不存在 (PID: {pid})")

    def resume_process(self, pid):
        """恢复进程"""
        try:
            psutil.Process(pid).resume()
        except psutil.NoSuchProcess:
            self.append_output(f"进程恢复失败，进程不存在 (PID: {pid})")

    def execute_cmd(self):
        """执行系统命令"""
        cmd = self.cmd_entry.text()
        disk = self.disk_combo.currentText()
        if not cmd:
            return
        
        try:
            result = subprocess.check_output(
                f"{disk} && {cmd}",
                shell=True,
                stderr=subprocess.STDOUT,
                text=True,
                timeout=10
            )
            self.append_output(f"$ {cmd}\n{result}")
        except subprocess.CalledProcessError as e:
            self.append_output(f"命令执行错误[{e.returncode}]:\n{e.output}")
        except subprocess.TimeoutExpired:
            self.append_output("命令执行超时（超过10秒）")

    def scan_disks(self):
        """扫描可用磁盘"""
        self.disk_combo.clear()
        try:
            disks = [partition.device for partition in psutil.disk_partitions()]
            self.disk_combo.addItems(disks)
        except Exception as e:
            self.append_output(f"磁盘扫描失败: {str(e)}")

    def append_output(self, text):
        """添加输出信息"""
        self.output_text.moveCursor(QTextCursor.MoveOperation.End)
        self.output_text.insertPlainText(f"[{time.strftime('%H:%M:%S')}] {text}\n")
        self.output_text.ensureCursorVisible()

    def closeEvent(self, event):
        self.stop_monitoring()
        super().closeEvent(event)

app = QApplication(sys.argv)
window = SystemMonitor()
window.show()
sys.exit(app.exec())